import os
import sys
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import matplotlib.pyplot as plt


g_scale = 2
g_resize = 100
g_count = 10
g_outsize = 10000

def heart_shape(t, scale):
    """ 根据参数t生成心形曲线上的点 """
    x = 16 * np.sin(t) ** 3
    y = -(13 * np.cos(t) - 5 * np.cos(2 * t) - 2 * np.cos(3 * t) - np.cos(4 * t))
    return (x * scale, y * scale)

def add_num(img, number):
    # 创建一个绘图对象
    draw = ImageDraw.Draw(img)
    # 设置字体和大小
    font = ImageFont.truetype("arial.ttf", 40)  # 你可以根据需要调整字体和大小

    # 获取图像的宽度和高度
    width, height = img.size

    # 计算数字的位置（右上角）
    _, _, text_width, text_height = draw.textbbox((0, 0), str(number), font=font)
    position = (width - text_width - 10, 10)  # 右上角位置，留一点边距

    # 绘制数字
    draw.text(position, str(number), fill=(255, 255, 255), font=font)  # 红色数字
    # return img

def place_images(image_files, output_size=(g_outsize, g_outsize)):
    print(f"len: {len(image_files)}")
    # 创建一个空白画布
    canvas = Image.new('RGB', output_size, 'white')
    
    # 加载图片并调整大小
    images = [Image.open(file).resize((g_resize, g_resize)) for file in image_files]
    
    # 计算心形曲线上的点
    points = [heart_shape(t, g_scale) for t in np.linspace(0, 2 * np.pi, len(images), endpoint=False)]
    # print(f"points: {points}")
    index = 0
    # 将图片放置到画布上
    for img, point in zip(images, points):
        x, y = point
        # 调整坐标使其位于画布中央
        x = int(output_size[0] / 2 + x - img.width / 2)
        y = int(output_size[1] / 2 + y - img.height / 2)
        add_num(img, index)
        # if index == 0:
        print(f"{index} -> {point}")
        canvas.paste(img, (x, y), img.convert('RGBA'))
        index += 1
    
    return canvas

def heart_test():
    # 假设图片文件存储在'pngs'文件夹中 99 张图片
    image_folder = './pngs'
    image_files = [os.path.join(image_folder, f) for f in os.listdir(image_folder) if f.endswith('.jpg') or f.endswith('.png')]
    
    # 生成心形排版
    heart_image = place_images(image_files, output_size=(g_outsize, g_outsize))
    
    # 保存结果
    heart_image.save('heart_layout.png')

# 九宫格
def cut_image_into_9(image_path, output_dir):
    # 打开图像
    image = Image.open(image_path)
    
    # 获取图像的宽度和高度
    width, height = image.size
    
    # 计算每个小图的宽度和高度
    grid_size = 3
    cell_width = width // grid_size
    cell_height = height // grid_size
    
    # 创建输出目录（如果不存在）
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 切割图像
    for row in range(grid_size):
        for col in range(grid_size):
            # 计算每个小图的边界
            left = col * cell_width
            upper = row * cell_height
            right = (col + 1) * cell_width
            lower = (row + 1) * cell_height
            
            # 裁剪小图
            cropped_image = image.crop((left, upper, right, lower))
            
            # 保存小图
            output_path = os.path.join(output_dir, f'part_{row}_{col}.png')
            cropped_image.save(output_path)
            print(f'Saved {output_path}')

def cut9():
    # 使用示例
    Image.MAX_IMAGE_PIXELS = 400000001  # 例如设置为10亿像素
    input_image_path = 'heart_layout.png'  # 输入图像路径
    output_directory = 'output_images'    # 输出图像目录

    cut_image_into_9(input_image_path, output_directory)

def run():
    # heart_test()
    cut9()

# prepare: pip install pillow
# python.exe heart.py 600 750 99 20000
if __name__ == '__main__':
    if len(sys.argv) > 1:
        g_scale = int(sys.argv[1])
        g_resize = int(sys.argv[2])
        g_count = int(sys.argv[3])
        g_outsize = int(sys.argv[4])
    print(f"scale: {g_scale}, resize: {g_resize}, count: {g_count}, outsize: {g_outsize}")
    run()